{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "9t3_pqAO4jjS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "#@title Load library\n",
        "TORCH_ENABLED = True\n",
        "SCIPY_ENABLED = True\n",
        "\n",
        "from transformers import pipeline\n",
        "import requests\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "from math import cos, sin\n",
        "import math as m\n",
        "import os\n",
        "import cv2\n",
        "import pandas as pd\n",
        "import gc\n",
        "try:\n",
        "  from scipy.spatial import cKDTree\n",
        "except:\n",
        "  SCIPY_ENABLED = False\n",
        "try:\n",
        "  import torch\n",
        "except:\n",
        "  TORCH_ENABLED = False"
      ],
      "metadata": {
        "id": "jAwwGXpiQ81b",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "#@title Define functions\n",
        "\n",
        "def clean():\n",
        "  for I in range(0,20):\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()\n",
        "\n",
        "def export_PLY(points, name_file=\"model.ply\", multiple_files=False, format_ascii=\"ascii\"):\n",
        "\n",
        "    # Write PLY header\n",
        "    header = [\n",
        "        \"ply\",\n",
        "        f\"format {(format_ascii)} 1.0\",\n",
        "        f\"element vertex {(points.shape[0])}\",\n",
        "        \"property float x\",\n",
        "        \"property float y\",\n",
        "        \"property float z\",\n",
        "        \"property uchar red\",\n",
        "        \"property uchar green\",\n",
        "        \"property uchar blue\",\n",
        "        \"end_header\\n\"\n",
        "    ]\n",
        "\n",
        "    if not(name_file.endswith(\".ply\")):\n",
        "        name_file+=\".ply\"\n",
        "\n",
        "    if \"ascii\" in format_ascii:\n",
        "        header_cons = \"\\n\".join(header)\n",
        "        if not(multiple_files):\n",
        "            with open(name_file, 'wb') as f:\n",
        "                np.savetxt(f, points, fmt='%g %g %g %d %d %d', header=header_cons, comments='')\n",
        "            with open(name_file, 'rb+') as fout:\n",
        "                fout.seek(-1, os.SEEK_END)\n",
        "                fout.truncate()\n",
        "        else:\n",
        "            def split_array(arr, n):\n",
        "                return [arr[i:i+n] for i in range(0, len(arr), n)]\n",
        "            for idx,points_ck in enumerate(split_array(points, 2097152)):\n",
        "                header = [\n",
        "                    \"ply\",\n",
        "                    f\"format {(format_ascii)} 1.0\",\n",
        "                    f\"element vertex {(points_ck.shape[0])}\",\n",
        "                    \"property float x\",\n",
        "                    \"property float y\",\n",
        "                    \"property float z\",\n",
        "                    \"property uchar red\",\n",
        "                    \"property uchar green\",\n",
        "                    \"property uchar blue\",\n",
        "                    \"end_header\\n\"\n",
        "                ]\n",
        "                name_file_t = name_file.split(\".ply\")[0]+\"_\"+str(idx)+\".ply\"\n",
        "                with open(name_file_t, 'wb') as f:\n",
        "                    np.savetxt(f, points_ck, fmt='%g %g %g %d %d %d', header=header_cons, comments='')\n",
        "                with open(name_file_t, 'rb+') as fout:\n",
        "                    fout.seek(-1, os.SEEK_END)\n",
        "                    fout.truncate()\n",
        "\n",
        "    else:\n",
        "        xyz = points[:, :3].astype(np.float32)\n",
        "        rgb = points[:, 3:6].astype(np.uint8)\n",
        "\n",
        "\n",
        "\n",
        "        if (multiple_files):\n",
        "            def split_array(arr, n):\n",
        "                return [arr[i:i+n] for i in range(0, len(arr), n)]\n",
        "            xyz = split_array(xyz, 2097152)\n",
        "            rgb = split_array(rgb, 2097152)\n",
        "        else:\n",
        "            xyz =[xyz]\n",
        "            rgb =[rgb]\n",
        "\n",
        "        for idx,xyz_ in enumerate(xyz):\n",
        "            header = [\n",
        "                \"ply\",\n",
        "                f\"format {(format_ascii)} 1.0\",\n",
        "                f\"element vertex {(xyz_.shape[0])}\",\n",
        "                \"property float x\",\n",
        "                \"property float y\",\n",
        "                \"property float z\",\n",
        "                \"property uchar red\",\n",
        "                \"property uchar green\",\n",
        "                \"property uchar blue\",\n",
        "                \"end_header\"\n",
        "            ]\n",
        "            structured_array = np.zeros(xyz_.shape[0], dtype=[\n",
        "                ('x', '<f4'),  # Little-endian float32\n",
        "                ('y', '<f4'),\n",
        "                ('z', '<f4'),\n",
        "                ('red', 'u1'),  # Unsigned byte (0-255)\n",
        "                ('green', 'u1'),\n",
        "                ('blue', 'u1')\n",
        "            ])\n",
        "            rgb_ = rgb[idx]\n",
        "            name_file_t = (name_file.split(\".ply\")[0]+\"_\"+str(idx)+\".ply\" if multiple_files else name_file)\n",
        "            structured_array['x'] = xyz_[:, 0]\n",
        "            structured_array['y'] = xyz_[:, 1]\n",
        "            structured_array['z'] = xyz_[:, 2]\n",
        "            structured_array['red'] = rgb_[:, 0]\n",
        "            structured_array['green'] = rgb_[:, 1]\n",
        "            structured_array['blue'] = rgb_[:, 2]\n",
        "            with open(name_file_t, 'wb') as f:\n",
        "                f.write('\\n'.join(header).encode('utf-8'))\n",
        "                f.write(b'\\n')  # Header ends with a newline\n",
        "                structured_array.tofile(f)\n",
        "\n",
        "def from2Dto3D_vectorized(arrayDepth, arrayPixel, maxZ, quality=1):\n",
        "    # Check if the input shapes are compatible\n",
        "    assert arrayPixel.shape[:2] == arrayDepth.shape, \"Pixel and Depth map shapes do not match\"\n",
        "    W, H = arrayDepth.shape  # Original dimensions\n",
        "\n",
        "    if quality > 1:\n",
        "        # Calculate new dimensions based on quality\n",
        "        new_W = (W - 1) * quality + 1\n",
        "        new_H = (H - 1) * quality + 1\n",
        "\n",
        "        # Generate new coordinates using linspace\n",
        "        new_X = np.linspace(0, W-1, new_W)\n",
        "        new_Y = np.linspace(0, H-1, new_H)\n",
        "\n",
        "        # Create meshgrid for new coordinates\n",
        "        xi, yi = np.meshgrid(new_X, new_Y, indexing='ij')  # shapes (new_W, new_H)\n",
        "\n",
        "        # Prepare indices for interpolation\n",
        "        x0 = np.floor(xi).astype(int)\n",
        "        x1 = np.clip(x0 + 1, 0, W-1)\n",
        "        y0 = np.floor(yi).astype(int)\n",
        "        y1 = np.clip(y0 + 1, 0, H-1)\n",
        "\n",
        "        dx = xi - x0\n",
        "        dy = yi - y0\n",
        "\n",
        "        # Interpolate arrayDepth\n",
        "        tl = arrayDepth[x0, y0]\n",
        "        tr = arrayDepth[x0, y1]\n",
        "        bl = arrayDepth[x1, y0]\n",
        "        br = arrayDepth[x1, y1]\n",
        "        interpolated_depth = (\n",
        "            (1 - dx) * (1 - dy) * tl +\n",
        "            dx * (1 - dy) * tr +\n",
        "            (1 - dx) * dy * bl +\n",
        "            dx * dy * br\n",
        "        )\n",
        "\n",
        "        # Interpolate arrayPixel for each channel\n",
        "        interpolated_pixel = np.zeros((new_W, new_H, 3), dtype=arrayPixel.dtype)\n",
        "        for c in range(3):\n",
        "            tl_c = arrayPixel[x0, y0, c]\n",
        "            tr_c = arrayPixel[x0, y1, c]\n",
        "            bl_c = arrayPixel[x1, y0, c]\n",
        "            br_c = arrayPixel[x1, y1, c]\n",
        "            interpolated_c = (\n",
        "                (1 - dx) * (1 - dy) * tl_c +\n",
        "                dx * (1 - dy) * tr_c +\n",
        "                (1 - dx) * dy * bl_c +\n",
        "                dx * dy * br_c\n",
        "            )\n",
        "            interpolated_pixel[:, :, c] = interpolated_c.astype(arrayPixel.dtype)\n",
        "\n",
        "        # Update variables to use interpolated arrays\n",
        "        W, H = new_W, new_H\n",
        "        arrayDepth = interpolated_depth\n",
        "        arrayPixel = interpolated_pixel\n",
        "        # Use new_X and new_Y for coordinates\n",
        "        final_X, final_Y = np.meshgrid(new_X, new_Y, indexing='ij')\n",
        "    else:\n",
        "        # Original coordinates\n",
        "        x = np.arange(W)\n",
        "        y = np.arange(H)\n",
        "        final_X, final_Y = np.meshgrid(x, y, indexing='ij')  # shapes (W, H)\n",
        "\n",
        "    # Calculate scaled Z values\n",
        "    d_min = np.min(arrayDepth)\n",
        "    d_max = np.max(arrayDepth)\n",
        "    if d_max - d_min > 1e-9:\n",
        "        Z = -(arrayDepth - d_min) / (d_max - d_min) * maxZ\n",
        "    else:\n",
        "        Z = np.zeros_like(arrayDepth, dtype=np.float64)\n",
        "\n",
        "    # Extract RGB components\n",
        "    R = arrayPixel[:, :, 0]\n",
        "    G = arrayPixel[:, :, 1]\n",
        "    B = arrayPixel[:, :, 2]\n",
        "\n",
        "    # Stack all components and reshape to Nx6\n",
        "    # Transpose to shape (H, W) for correct ordering when using 'ij' indexing\n",
        "    points = np.stack([\n",
        "        final_Y.T, final_X.T, Z.T,\n",
        "        R.T, G.T, B.T\n",
        "    ], axis=-1)\n",
        "    points = points.reshape(-1, 6)\n",
        "\n",
        "    return points\n",
        "\n",
        "\n",
        "def from2Dto3D_vectorized_torch(arrayDepth, arrayPixel, maxZ, quality=1):\n",
        "    # Ensure inputs are PyTorch tensors\n",
        "    assert isinstance(arrayDepth, torch.Tensor), \"arrayDepth must be a torch.Tensor\"\n",
        "    assert isinstance(arrayPixel, torch.Tensor), \"arrayPixel must be a torch.Tensor\"\n",
        "\n",
        "    # Check shape compatibility\n",
        "\n",
        "    assert arrayPixel.shape[:2] == arrayDepth.shape, \"Pixel and Depth map shapes do not match\"\n",
        "    W, H = arrayDepth.shape  # Original dimensions\n",
        "\n",
        "    if quality > 1:\n",
        "        # Calculate new dimensions\n",
        "        new_W = (W - 1) * quality + 1\n",
        "        new_H = (H - 1) * quality + 1\n",
        "\n",
        "        # Generate coordinates (CHANGED: torch.linspace instead of np.linspace)\n",
        "        new_X = torch.linspace(0, W-1, new_W, device=arrayDepth.device)\n",
        "        new_Y = torch.linspace(0, H-1, new_H, device=arrayDepth.device)\n",
        "\n",
        "        # Create meshgrid (CHANGED: torch.meshgrid with indexing='ij')\n",
        "        xi, yi = torch.meshgrid(new_X, new_Y, indexing='ij')\n",
        "\n",
        "        # Prepare indices (CHANGED: torch.floor/int/clamp)\n",
        "        x0 = torch.floor(xi).long()\n",
        "        x1 = torch.clamp(x0 + 1, 0, W-1)\n",
        "        y0 = torch.floor(yi).long()\n",
        "        y1 = torch.clamp(y0 + 1, 0, H-1)\n",
        "\n",
        "        dx = xi - x0\n",
        "        dy = yi - y0\n",
        "\n",
        "        # Interpolate depth (CHANGED: Tensor indexing)\n",
        "        tl = arrayDepth[x0, y0]\n",
        "        tr = arrayDepth[x0, y1]\n",
        "        bl = arrayDepth[x1, y0]\n",
        "        br = arrayDepth[x1, y1]\n",
        "        interpolated_depth = (\n",
        "            (1 - dx) * (1 - dy) * tl +\n",
        "            dx * (1 - dy) * tr +\n",
        "            (1 - dx) * dy * bl +\n",
        "            dx * dy * br\n",
        "        )\n",
        "\n",
        "        # Interpolate pixel (CHANGED: Tensor operations preserve dtype)\n",
        "        interpolated_pixel = torch.zeros((new_W, new_H, 3),\n",
        "                                      dtype=arrayPixel.dtype,\n",
        "                                      device=arrayPixel.device)\n",
        "        for c in range(3):\n",
        "            tl_c = arrayPixel[x0, y0, c]\n",
        "            tr_c = arrayPixel[x0, y1, c]\n",
        "            bl_c = arrayPixel[x1, y0, c]\n",
        "            br_c = arrayPixel[x1, y1, c]\n",
        "            interpolated_c = (\n",
        "                (1 - dx) * (1 - dy) * tl_c +\n",
        "                dx * (1 - dy) * tr_c +\n",
        "                (1 - dx) * dy * bl_c +\n",
        "                dx * dy * br_c\n",
        "            )\n",
        "            interpolated_pixel[:, :, c] = interpolated_c.to(arrayPixel.dtype)\n",
        "\n",
        "        W, H = new_W, new_H\n",
        "        arrayDepth = interpolated_depth\n",
        "        arrayPixel = interpolated_pixel\n",
        "        final_X, final_Y = xi, yi\n",
        "    else:\n",
        "        # Original coordinates (CHANGED: torch.arange)\n",
        "        x = torch.arange(W, device=arrayDepth.device)\n",
        "        y = torch.arange(H, device=arrayDepth.device)\n",
        "        final_X, final_Y = torch.meshgrid(x, y, indexing='ij')\n",
        "\n",
        "    # Calculate Z (CHANGED: Tensor operations)\n",
        "    d_min = torch.min(arrayDepth)\n",
        "    d_max = torch.max(arrayDepth)\n",
        "    if (d_max - d_min).item() > 1e-9:  # CHANGED: .item() for scalar comparison\n",
        "        Z = -(arrayDepth - d_min) / (d_max - d_min) * maxZ\n",
        "    else:\n",
        "        Z = torch.zeros_like(arrayDepth)\n",
        "\n",
        "    # Extract RGB (CHANGED: Tensor slicing)\n",
        "    R = arrayPixel[:, :, 0]\n",
        "    G = arrayPixel[:, :, 1]\n",
        "    B = arrayPixel[:, :, 2]\n",
        "\n",
        "    # Stack components (CHANGED: permute instead of transpose)\n",
        "    points = torch.stack([\n",
        "        final_Y.permute(1, 0),  # Equivalent to .T in NumPy\n",
        "        final_X.permute(1, 0),\n",
        "        Z.permute(1, 0),\n",
        "        R.permute(1, 0),\n",
        "        G.permute(1, 0),\n",
        "        B.permute(1, 0)\n",
        "    ], dim=-1)\n",
        "\n",
        "    return points.view(-1, 6)  # CHANGED: view instead of reshape\n",
        "\n",
        "def transform_points(points, translate=(0, 0, 0), rotate=(0, 0, 0), scale=(1,1,1)):\n",
        "    \"\"\"Apply 3D transformations to points (rotation first, then translation)\"\"\"\n",
        "    # Convert rotation angles to radians if needed (modify if using degrees)\n",
        "    rx, ry, rz = rotate\n",
        "\n",
        "\n",
        "    # Create rotation matrices\n",
        "    # X-axis rotation\n",
        "    rot_x = np.array([\n",
        "        [1, 0, 0],\n",
        "        [0, cos(rx), -sin(rx)],\n",
        "        [0, sin(rx), cos(rx)]\n",
        "    ], dtype=np.float64)\n",
        "\n",
        "    # Y-axis rotation\n",
        "    rot_y = np.array([\n",
        "        [cos(ry), 0, sin(ry)],\n",
        "        [0, 1, 0],\n",
        "        [-sin(ry), 0, cos(ry)]\n",
        "    ], dtype=np.float64)\n",
        "\n",
        "    # Z-axis rotation\n",
        "    rot_z = np.array([\n",
        "        [cos(rz), -sin(rz), 0],\n",
        "        [sin(rz), cos(rz), 0],\n",
        "        [0, 0, 1]\n",
        "    ], dtype=np.float64)\n",
        "\n",
        "\n",
        "    # Combined rotation matrix (Z-Y-X order)\n",
        "    rotation_matrix = rot_z @ rot_y @ rot_x\n",
        "\n",
        "    points[:,0:3] = points[:,0:3].astype(np.float64)\n",
        "\n",
        "    points[:,0:3] = (points[:,0:3] * np.array(scale, dtype=np.float64)).astype(np.float64)\n",
        "\n",
        "    # Apply rotation\n",
        "    points[:,0:3] = (points[:,0:3] @ rotation_matrix.T).astype(np.float64)\n",
        "\n",
        "    # Apply translation\n",
        "    points[:,0:3] = (points[:,0:3] + np.array(translate, dtype=np.float64)).astype(np.float64)\n",
        "\n",
        "    return points\n",
        "\n",
        "def project_points(points, fov=60, rotation=None, translation=None, scale=None, aspect_ratio=1, img_size=(1000,1000)):\n",
        "    \"\"\"Fast 3D->2D projection with perspective correction\"\"\"\n",
        "    if rotation is None:\n",
        "        rotation = np.zeros(3)\n",
        "    if translation is None:\n",
        "        translation = np.zeros(3)\n",
        "    if scale is None:\n",
        "        scale = np.ones(3)\n",
        "\n",
        "    # Calcolo parametri camera\n",
        "    fov_rad = m.radians(fov)\n",
        "    focal_length = img_size[1] / (2 * m.tan(fov_rad / 2))  # Focal length verticale\n",
        "\n",
        "    # Trasformazioni (camera guarda lungo -Z)\n",
        "\n",
        "    points = transform_points(points, translate=translation, rotate=rotation, scale=scale)\n",
        "\n",
        "    # Seleziona punti DAVANTI alla camera (z < 0)\n",
        "    valid = points[:, 2] < 0\n",
        "    points_t = points #[valid]\n",
        "\n",
        "    if len(points_t) == 0:\n",
        "        return np.empty((0, 2)), np.array([]), np.empty((0, 3))\n",
        "\n",
        "    # Calcolo coordinate proiettate\n",
        "    z = -points_t[:, 2]  # Converti in distanza positiva\n",
        "    x_proj = (points_t[:, 0] * focal_length) / (z * aspect_ratio)\n",
        "    y_proj = (points_t[:, 1] * focal_length) / z\n",
        "\n",
        "    # Conversione a coordinate immagine (Y non invertito)\n",
        "    pixel_x = (x_proj + img_size[0]/2).astype(int)\n",
        "    pixel_y = (y_proj + img_size[1]/2).astype(int)  # Rimosso il segno negativo\n",
        "\n",
        "    # Clip e normalizzazione\n",
        "    pixel_x = np.clip(pixel_x, 0, img_size[0]-1)\n",
        "    pixel_y = np.clip(pixel_y, 0, img_size[1]-1)\n",
        "    z_norm = (z - z.min()) / (z.max() - z.min())\n",
        "\n",
        "    return np.column_stack([pixel_x, pixel_y]), z_norm, points_t[:, 3:6]\n",
        "\n",
        "def project_points_ortho(points, rotation=None, translation=None, scale=None):\n",
        "    \"\"\"Orthographic 3D->2D projection\"\"\"\n",
        "    if rotation is None:\n",
        "        rotation = np.zeros(3)\n",
        "    if translation is None:\n",
        "        translation = np.zeros(3)\n",
        "    if scale is None:\n",
        "        scale = np.ones(3)\n",
        "\n",
        "\n",
        "    # Apply transformations\n",
        "    points = transform_points(points, translate=translation, rotate=rotation, scale=scale)\n",
        "\n",
        "    # Orthographic projection\n",
        "    x_proj = points[:, 0]  # Simple scaling + centering offset\n",
        "    y_proj = points[:, 1]\n",
        "\n",
        "    # Normalize depth (similar to perspective version)\n",
        "    z = points[:, 2]\n",
        "    z_min = z.min()\n",
        "    z_max = z.max()\n",
        "    z_range = z_max - z_min + 1e-5  # Avoid division by zero\n",
        "    z_norm = (z - z_min) / z_range\n",
        "\n",
        "    colors = points[:,3:6]\n",
        "\n",
        "    return np.column_stack([x_proj, y_proj]), z_norm, colors\n",
        "\n",
        "def render_points_fast(points, rotation=None, translation=None, scale=None,\n",
        "                      img_size=(1000, 1000), color=False, cameraType=0, fov=60, correction=False, ksize=(2,2), thresh=2):\n",
        "    \"\"\"Ultra-fast rendering using pure NumPy and PIL\"\"\"\n",
        "    # Project points to 2D\n",
        "    if cameraType==0:\n",
        "        proj, depth, colors = project_points_ortho(points, rotation=rotation, translation=translation)\n",
        "    else:\n",
        "        proj, depth, colors = project_points(points, fov=fov, aspect_ratio=img_size[1]/img_size[1], rotation=rotation, translation=translation, scale=scale, img_size=img_size)\n",
        "\n",
        "    if isinstance(proj, torch.Tensor):\n",
        "      proj=proj.detach().cpu().numpy()\n",
        "      depth=depth.detach().cpu().numpy()\n",
        "      colors=colors.detach().cpu().numpy()\n",
        "\n",
        "    # Convert to integer coordinates and filter valid points\n",
        "    proj = proj.astype(np.int32)\n",
        "    valid = (proj[:, 0] >= 0) & (proj[:, 0] < img_size[0]) & \\\n",
        "            (proj[:, 1] >= 0) & (proj[:, 1] < img_size[1])\n",
        "\n",
        "    proj = proj[valid]\n",
        "    depth = depth[valid]\n",
        "    colors = colors[valid]\n",
        "\n",
        "    # Grayscale rendering based on depth\n",
        "    max_depth = depth.max() if depth.size > 0 else 1\n",
        "    min_depth = depth.min() if depth.size > 0 else 0\n",
        "    if max_depth - min_depth > 0:\n",
        "        depth_normalized = (depth - min_depth) / (max_depth - min_depth)\n",
        "    else:\n",
        "        depth_normalized = np.zeros_like(depth)\n",
        "    intensities = (255 - (depth_normalized * 255)).astype(np.uint16)\n",
        "\n",
        "    # Sort points by descending intensity to prioritize closer points\n",
        "    sorted_indices = np.argsort(-intensities)\n",
        "    proj_sorted = proj[sorted_indices]\n",
        "    y_coords = proj_sorted[:, 1]\n",
        "    x_coords = proj_sorted[:, 0]\n",
        "\n",
        "    df = pd.DataFrame({'y': y_coords, 'x': x_coords})\n",
        "    unique_indices = df.drop_duplicates().index.to_numpy()\n",
        "\n",
        "    # Extract unique coordinates and colors\n",
        "    unique_y = y_coords[unique_indices]\n",
        "    unique_x = x_coords[unique_indices]\n",
        "\n",
        "    if color:\n",
        "        # Extract colors from valid points (assuming points are Nx[x,y,z,r,g,b])\n",
        "        point_colors = colors.astype(np.uint16)\n",
        "        colors_sorted = point_colors[sorted_indices]\n",
        "\n",
        "        unique_colors = colors_sorted[unique_indices]\n",
        "\n",
        "        # Create color buffer and assign colors\n",
        "        color_buffer = np.ones((img_size[1], img_size[0], 3))\n",
        "        color_buffer = color_buffer*-1\n",
        "        color_buffer[unique_y, unique_x] = unique_colors\n",
        "\n",
        "        img = color_buffer\n",
        "    else:\n",
        "        point_colors = intensities\n",
        "        colors_sorted = point_colors[sorted_indices]\n",
        "\n",
        "        unique_colors = colors_sorted[unique_indices]\n",
        "\n",
        "        # Create color buffer and assign colors\n",
        "        color_buffer = np.ones((img_size[1], img_size[0]))\n",
        "        color_buffer = color_buffer*-1\n",
        "        color_buffer[unique_y, unique_x] = unique_colors\n",
        "\n",
        "        img = color_buffer\n",
        "\n",
        "\n",
        "    if correction:\n",
        "        point_colors = intensities\n",
        "        colors_sorted = point_colors[sorted_indices]\n",
        "\n",
        "        unique_colors = colors_sorted[unique_indices]\n",
        "\n",
        "        # Create color buffer and assign colors\n",
        "        depth_buffer = np.zeros((img_size[1], img_size[0]))\n",
        "        depth_buffer[unique_y, unique_x] = unique_colors\n",
        "\n",
        "        img_depth = depth_buffer\n",
        "\n",
        "        img = np.array(img)\n",
        "        img_dep = np.array(img_depth)\n",
        "\n",
        "        img_blur = blur_image_excluding_black(img, ksize, thresh)\n",
        "        #print(np.unique(img_blur))\n",
        "        mask = img_dep < thresh\n",
        "\n",
        "        (img[mask]) = (img_blur[mask])\n",
        "\n",
        "        img = img\n",
        "\n",
        "    return img\n",
        "\n",
        "def clean_points(points, k, m):\n",
        "    points = np.asarray(points)\n",
        "    if points.size == 0:\n",
        "        return points.copy()\n",
        "    if k <= 0:\n",
        "        raise ValueError(\"k must be a positive integer\")\n",
        "    n = len(points)\n",
        "    if k > n - 1:\n",
        "        return np.empty((0, points.shape[1]))\n",
        "    tree = cKDTree(points, balanced_tree=False)\n",
        "    # Query k+1 to include the point itself, then select the k-th neighbor\n",
        "    distances, _ = tree.query(points, k=k+1, workers=-1)\n",
        "    kth_distances = distances[:, k]  # k-th neighbor after excluding self\n",
        "    mask = kth_distances <= m\n",
        "    return points[mask]\n",
        "\n",
        "def interpolate_points(points, alpha=0.5, n_c=3):\n",
        "    \"\"\"\n",
        "    Interpolate three new points for each original point, positioned between the original and its three nearest neighbors.\n",
        "\n",
        "    Parameters:\n",
        "    points (numpy.ndarray): Input array of shape (N, 6) where each row is (x, y, z, r, g, b).\n",
        "    alpha (float): Interpolation factor (0.0 = original point, 1.0 = neighbor). Default is 0.5 (midpoint).\n",
        "\n",
        "    Returns:\n",
        "    numpy.ndarray: Array of interpolated points with shape (3*N, 6).\n",
        "    \"\"\"\n",
        "    # Extract coordinates and colors\n",
        "    coords = points[:, :3]\n",
        "    colors = points[:, 3:]\n",
        "    N = coords.shape[0]\n",
        "\n",
        "    # Build KDTree for efficient neighbor lookup\n",
        "    tree = cKDTree(coords, balanced_tree=False)\n",
        "\n",
        "    # Query for 4 nearest neighbors (including self), then exclude self\n",
        "    _, indices = tree.query(coords, k=n_c+1, workers=-1)\n",
        "    neighbor_indices = indices[:, 1:(n_c+1)]  # Shape (N, 3)\n",
        "\n",
        "    # Prepare indices for vectorized operations\n",
        "    original_indices = np.repeat(np.arange(N), n_c)\n",
        "    neighbors_flat = neighbor_indices.ravel()\n",
        "\n",
        "    # Gather original and neighbor data\n",
        "    original_coords = coords[original_indices]\n",
        "    neighbor_coords = coords[neighbors_flat]\n",
        "\n",
        "    original_colors = colors[original_indices]\n",
        "    neighbor_colors = colors[neighbors_flat]\n",
        "\n",
        "    # Interpolate coordinates and colors\n",
        "    interpolated_coords = (1 - alpha) * original_coords + alpha * neighbor_coords\n",
        "    interpolated_colors = (1 - alpha) * original_colors + alpha * neighbor_colors\n",
        "\n",
        "    # Combine into new points array\n",
        "    new_points = np.hstack((interpolated_coords, interpolated_colors))\n",
        "\n",
        "    combined_points = np.vstack((points, new_points))\n",
        "\n",
        "    return combined_points\n",
        "\n",
        "def convolve2d_np(image, kernel, mode='same'):\n",
        "    # Flip the kernel for convolution\n",
        "    kernel = np.flipud(np.fliplr(kernel))\n",
        "    k_h, k_w = kernel.shape\n",
        "    i_h, i_w = image.shape\n",
        "\n",
        "    # Determine padding\n",
        "    if mode == 'same':\n",
        "        pad_top = (k_h - 1) // 2\n",
        "        pad_bottom = (k_h - 1) - pad_top\n",
        "        pad_left = (k_w - 1) // 2\n",
        "        pad_right = (k_w - 1) - pad_left\n",
        "        padded_image = np.pad(image, ((pad_top, pad_bottom), (pad_left, pad_right)), mode='constant')\n",
        "    elif mode == 'valid':\n",
        "        padded_image = image\n",
        "    elif mode == 'full':\n",
        "        padded_image = np.pad(image, ((k_h-1, k_h-1), (k_w-1, k_w-1)), mode='constant')\n",
        "    else:\n",
        "        raise ValueError(\"Mode must be 'same', 'valid', or 'full'\")\n",
        "\n",
        "    # Generate sliding windows\n",
        "    windows = np.lib.stride_tricks.sliding_window_view(padded_image, (k_h, k_w))\n",
        "\n",
        "    # Perform convolution by summing element-wise multiplication\n",
        "    result = np.sum(windows * kernel.reshape(1, 1, k_h, k_w), axis=(-2, -1))\n",
        "\n",
        "    return result\n",
        "\n",
        "def blur_image_excluding_black(image, kernel_size=(4,4), threshold=2):\n",
        "    # Create mask for non-black pixels (all channels zero)\n",
        "    if image.ndim == 3:\n",
        "        mask = np.any(image > threshold, axis=-1).astype(float)\n",
        "    else:\n",
        "        mask = (image > threshold).astype(float)\n",
        "\n",
        "    kernel = np.ones(kernel_size)\n",
        "\n",
        "    if image.ndim == 3:\n",
        "        blurred = np.ones_like(image)*-1\n",
        "        for c in range(image.shape[2]):\n",
        "            channel = image[:, :, c]\n",
        "            masked_channel = channel * mask\n",
        "            sum_matrix = convolve2d_np(masked_channel, kernel, mode='same')\n",
        "            count_matrix = convolve2d_np(mask, kernel, mode='same')\n",
        "            with np.errstate(divide='ignore', invalid='ignore'):\n",
        "                mean_matrix = sum_matrix / count_matrix\n",
        "                # Where count is zero, use original pixel if non-black, else 0\n",
        "                blurred_channel = np.where(count_matrix > 1, mean_matrix, -1)\n",
        "            blurred[:, :, c] = blurred_channel\n",
        "    else:\n",
        "        masked_image = image * mask\n",
        "        sum_matrix = convolve2d_np(masked_image, kernel, mode='same')\n",
        "        count_matrix = convolve2d_np(mask, kernel, mode='same')\n",
        "        with np.errstate(divide='ignore', invalid='ignore'):\n",
        "            mean_matrix = sum_matrix / count_matrix\n",
        "            blurred = np.where(count_matrix > 0, mean_matrix, masked_image)\n",
        "\n",
        "    return blurred"
      ],
      "metadata": {
        "id": "PoTzLauTQlHl",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "from google.colab import files\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "import PIL\n",
        "from PIL import Image, ImageDraw, ImageFilter\n",
        "\n",
        "#@title #Load image from upload file\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "loadImage = False\n",
        "mask = Image.new('RGBA', (512, 512),  (255, 255, 255))\n",
        "initt = None\n",
        "url = None\n",
        "\n",
        "\n",
        "display(\"Load Image:\")\n",
        "Up = files.upload()\n",
        "for k, v in Up.items():\n",
        "    initt=Image.open(k)\n",
        "\n",
        "\n",
        "print(initt.size)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "4kuU5vFnlfGY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6aEoYK7UPbPR",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Generate depth map\n",
        "test =True #@param {\"type\":\"boolean\"}\n",
        "# load pipe\n",
        "pipe = pipeline(task=\"depth-estimation\", model=\"depth-anything/Depth-Anything-V2-Large-hf\")\n",
        "\n",
        "# load image\n",
        "if test:\n",
        " url = 'https://raw.githubusercontent.com/chri002/ComfyUI_depthMapOperation/refs/heads/main/assets/start.jpg'\n",
        " image = Image.open(requests.get(url, stream=True).raw)\n",
        " image = image.resize((int(image.size[0]), int(image.size[1])))\n",
        " print(image.size)\n",
        "else:\n",
        " image=initt.convert(\"RGB\")\n",
        "# inference\n",
        "depth_img = pipe(image)[\"depth\"]"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Generate Cloud Points\n",
        "depth =512#@param {\"type\" : \"number\"}\n",
        "\n",
        "quality=2 #@param {\"type\" : \"integer\"}\n",
        "img = torch.from_numpy((np.array(image)))\n",
        "depth_map = torch.from_numpy(np.array(depth_img))\n",
        "points = (from2Dto3D_vectorized_torch(depth_map, img, depth, quality)).cpu().numpy()"
      ],
      "metadata": {
        "id": "u9S_ulEFRmiI",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Clean Cloud Points\n",
        "\n",
        "k = 20#@param {\"type\" : \"integer\"}\n",
        "distance = 16#@param {\"type\" : \"number\"}\n",
        "\n",
        "points = transform_points(points, translate=(0, 0, 0), rotate=(0, 0, 0), scale=(1,1,2))\n",
        "points = (clean_points(points, k=k, m=distance))\n",
        "points = transform_points(points, translate=(0, 0, 0), rotate=(0, 0, 0), scale=(1,1,0.5))"
      ],
      "metadata": {
        "id": "gEjBciYYldPF",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Local Rotation\n",
        "rx=0#@param {\"type\" : \"number\"}\n",
        "ry=0#@param {\"type\" : \"number\"}\n",
        "rz=0#@param {\"type\" : \"number\"}\n",
        "\n",
        "mx, Mx, my, My, mz, Mz = points[:,0].min(), points[:,0].max(), points[:,1].min(), points[:,1].max(), points[:,2].min(), points[:,2].max()\n",
        "tx = -(Mx+mx)/2\n",
        "ty = -(My+my)/2\n",
        "tz = -(Mz+mz)/2\n",
        "\n",
        "points = transform_points(points, translate=(tx, ty, tz), rotate=(0, 0, 0), scale=(1,1,1))\n",
        "points = transform_points(points, translate=(0, 0, 0), rotate=(rx/180*m.pi, ry/180*m.pi, rz/180*m.pi), scale=(1,1,1))\n",
        "points = transform_points(points, translate=(-tx, -ty, -tz), rotate=(0, 0, 0), scale=(1,1,1))"
      ],
      "metadata": {
        "id": "pV_6QxXTSfF7",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Translation\n",
        "tx=0#@param {\"type\" : \"number\"}\n",
        "ty=0#@param {\"type\" : \"number\"}\n",
        "tz=0#@param {\"type\" : \"number\"}\n",
        "\n",
        "points = transform_points(points, translate=(tx, ty, tz), rotate=(0, 0, 0), scale=(1,1,1))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "2HcdCFzirZit"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Scale\n",
        "sx=1#@param {\"type\" : \"number\"}\n",
        "sy=1#@param {\"type\" : \"number\"}\n",
        "sz=1#@param {\"type\" : \"number\"}\n",
        "\n",
        "points = transform_points(points, translate=(0,0,0), rotate=(0, 0, 0), scale=(sx,sy,sz))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "8Xp1q3DzrONh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Points interpolation\n",
        "interpolate = 3 #@param {\"type\":\"number\"}\n",
        "points =  (interpolate_points(points, alpha=0.5, n_c=interpolate))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "uQXuFh6Gr30U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Generate image\n",
        "import datetime\n",
        "start = datetime.datetime.now()\n",
        "\n",
        "tend = 0\n",
        "\n",
        "\n",
        "\n",
        "width,height = image.size\n",
        "translation = np.array([-width/2,-height/2,0])\n",
        "scale = np.array([1,1,-1])\n",
        "fov=45 #@param {\"type\" : \"number\"}\n",
        "correction = True #@param{\"type\":\"boolean\"}\n",
        "k_size =2 #@param{\"type\" : \"number\"}\n",
        "points_cop = transform_points(points.copy(), translate=translation, scale=scale)\n",
        "translation = np.array([0,0, -(height/m.sin(m.radians(fov)))])\n",
        "points_cop = transform_points(points_cop, translate=translation)\n",
        "\n",
        "\n",
        "end_img = (render_points_fast(points_cop, img_size=(width,height), color=True, cameraType=1, fov=fov, correction=correction, ksize=(k_size,k_size), thresh=3))\n",
        "\n",
        "background = 125*np.ones((int(image.size[0]/20),int(image.size[1]/20),3))\n",
        "background[1::2, ::2,:]=255\n",
        "background[::2, 1::2,:]=255\n",
        "background= cv2.resize(background, image.size, interpolation = cv2.INTER_NEAREST)\n",
        "\n",
        "mask = end_img[:,:,0]==-1\n",
        "end_img[mask] = background[mask]\n",
        "mask = end_img[:,:,1]==-1\n",
        "end_img[mask] = background[mask]\n",
        "mask = end_img[:,:,2]==-1\n",
        "end_img[mask] = background[mask]\n",
        "\n",
        "tend=datetime.datetime.now()\n",
        "\n",
        "delta=tend-start\n",
        "\n",
        "print(\"vertices: \",points.shape[0],\"\\ntime: \",delta)\n",
        "\n",
        "img_end = Image.fromarray(end_img.astype(np.uint8))\n",
        "\n",
        "images = [image, img_end]\n",
        "widths, heights = zip(*(i.size for i in images))\n",
        "\n",
        "total_width = sum(widths)\n",
        "max_height = max(heights)\n",
        "\n",
        "new_im = Image.new('RGB', (total_width, max_height))\n",
        "\n",
        "x_offset = 0\n",
        "for im in images:\n",
        "  new_im.paste(im, (x_offset,0))\n",
        "  x_offset += im.size[0]\n",
        "\n",
        "display(new_im)"
      ],
      "metadata": {
        "id": "f70devzISk8j",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Show Image\n",
        "img_end"
      ],
      "metadata": {
        "id": "N56fP1HRrGCD",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Export PLY\n",
        "export_PLY(points.detach().cpu().numpy(), name_file=\"model.ply\", multiple_files=False, format_ascii=\"binary_little_endian\")"
      ],
      "metadata": {
        "id": "osWtc_Ln0T67",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Clean memory\n",
        "clean()"
      ],
      "metadata": {
        "id": "B7WfeVFanOqx",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}